#pragma once
#include <iostream>
#include <list>
#include <vector>
#include <algorithm>
#include <array>
#include <time.h>
#include <queue>
#include <stack>
#include <string>
#include <set>
#include <map>
#include <functional>
#include <cassert>
using namespace std;

// 颜色枚举
enum Colour
{
	RED,
	BLACK
};

// 结点类
template <class K, class V>
struct RBTreeNode
{
	RBTreeNode<K, V> *_left;
	RBTreeNode<K, V> *_right;
	RBTreeNode<K, V> *_parent;

	pair<K, V> _kv;
	Colour _col;

	RBTreeNode(const pair<K, V> &kv)
		: _left(nullptr), _right(nullptr), _parent(nullptr), _kv(kv), _col(RED)
	{
	}
};

// RBTree类
template <class K, class V>
struct RBTree
{
	typedef RBTreeNode<K, V> Node;

public:
	int _rotateCount = 0;

public:
	// 插入创建红黑树
	bool Insert(const pair<K, V> &kv)
	{
		// 1.二叉搜索树插入

		// 树为空 插入结点即为根
		if (_root == nullptr)
		{
			_root = new Node(kv);
			_root->_col = BLACK; // 根节点颜色为黑
			return true;
		}

		// 不为空 定位至合适位置
		Node *dad = nullptr;
		Node *cur = _root;
		while (cur)
		{
			if (cur->_kv.first < kv.first)
			{
				dad = cur;
				cur = cur->_right;
			}
			else if (cur->_kv.first > kv.first)
			{
				dad = cur;
				cur = cur->_left;
			}
			else
			{
				return false;
			}
		}

		// 父连接子
		cur = new Node(kv);
		if (dad->_kv.first < kv.first)
		{
			dad->_right = cur;
		}
		else
		{
			dad->_left = cur;
		}
		// 子继承父
		cur->_parent = dad;

		// 2.构建红黑树

		// cur[红]此时定位新插入位置 若父存在且为红进入while
		while (dad && dad->_col == RED)
		{
			// 之后的循环来到此处 判断dad && dad->_col == RED
			// 父为空 则cur为根 [最后控制cur变黑] 循环结束; 父为黑 无需操作
			Node *grandpa = dad->_parent;

			// 若父为红 祖父定存在且为黑
			assert(grandpa && grandpa->_col == BLACK);

			// 2.1 判断uncle位置 分类一：父为左子树
			if (dad == grandpa->_left)
			{
				// uncle为右子树
				Node *uncle = grandpa->_right;

				// 情况一、uncle存在为红 变色+回溯
				// dad-uncle变黑 grandpa变红
				if (uncle && uncle->_col == RED)
				{
					// 变色
					dad->_col = uncle->_col = BLACK;
					grandpa->_col = RED;
					// 回溯
					cur = grandpa;
					dad = cur->_parent;
				}

				// 情况二-三：uncle不存在 或 存在且为黑
				else
				{
					if (cur == dad->_left)
					{
						// 情况二：右单旋+变色
						RotateR(grandpa);
						dad->_col = BLACK;
						grandpa->_col = RED;
					}
					else
					{
						// 情况三：左右双旋+变色
						RotateL(dad);
						RotateR(grandpa);
						cur->_col = BLACK;
						grandpa->_col = RED;
					}

					break;
				}
			}
			// 2.2 判断uncle位置 父为右子树
			else
			{
				// uncle为左子树
				Node *uncle = grandpa->_left;
				// 情况一
				if (uncle && uncle->_col == RED)
				{
					// 变色
					dad->_col = uncle->_col = BLACK;
					grandpa->_col = RED;
					// 回溯
					cur = grandpa;
					dad = cur->_parent;
				}
				else
				{
					if (cur == dad->_right)
					{
						// 情况二：左单旋+变色
						RotateL(grandpa);
						dad->_col = BLACK;
						grandpa->_col = RED;
					}
					else
					{
						// 情况三：右左单旋+变色
						RotateR(dad);
						RotateL(grandpa);
						cur->_col = BLACK;
						grandpa->_col = RED;
					}

					break;
				}
			}
		}

		// 父不存在的情况：即该树是空树 插入的结点就是根 在该函数调用初就判断了直接把该节点变为黑即可
		// cur的父为黑 则不用操作
		_root->_col = BLACK;
		return true;
	}

	// 中序遍历
	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

	// 判断是否是红黑树
	bool IsRBTree()
	{
		// 空树默认为真
		if (_root == nullptr)
		{
			return true;
		}

		// 根节点颜色
		if (_root->_col == RED)
		{
			cout << "根结点颜色错误!" << endl;
			return false;
		}

		// 某一条路径黑色结点数量
		int certain = 0;
		Node *cur = _root;
		while (cur)
		{
			if (cur->_col == BLACK)
				++certain;
			cur = cur->_left;
		}

		int BNcount = 0;
		return PrevJudge(_root, BNcount, certain);
	}

	Node *Find(const K &key)
	{
		Node *cur = _root;
		while (cur)
		{
			if (cur->_kv.first < key)
			{
				cur = cur->_right;
			}
			else if (cur->_kv.first > key)
			{
				cur = cur->_left;
			}
			else
			{
				return cur;
			}
		}

		return nullptr;
	}

	// 获取树的高度
	int Height()
	{
		return _Height(_root);
	}

	// 获取树的节点数
	size_t Size()
	{
		return _Size(_root);
	}

	~RBTree()
	{
		cout << "析构函数调用" << endl;
		_Destroy(_root);
		_root = nullptr;
	}

	void _Destroy(Node *root)
	{
		if (root == nullptr)
		{
			return;
		}

		_Destroy(root->_left);
		_Destroy(root->_right);
		delete root;
	}

private:
	// BNcount: black node count
	// certain: 某一条路径黑色结点数量
	bool PrevJudge(Node *root, int BNcount, int &certain)
	{
		// 遇空 -- 当前路径结束
		if (root == nullptr)
		{
			// 判断性质4
			if (BNcount != certain)
			{
				cout << "性质4违例: 存在某一路径黑色节点的数量不等!" << endl;
				return false;
			}
			return true;
		}

		// 遇黑--值++
		if (root->_col == BLACK)
		{
			++BNcount;
		}

		// 遇红--连坐判断
		if (root->_col == RED && root->_parent->_col == RED)
		{
			cout << "性质3违例: 存在连续红节点!" << endl;
			return false;
		}

		return PrevJudge(root->_left, BNcount, certain) && PrevJudge(root->_right, BNcount, certain);
	}

	// 中序遍历
	void _InOrder(Node *root)
	{
		if (root == nullptr)
		{
			return;
		}

		_InOrder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << endl;
		_InOrder(root->_right);
	}

	// 获取树的节点个数
	size_t _Size(Node *root)
	{
		if (root == NULL)
			return 0;

		return _Size(root->_left) + _Size(root->_right) + 1;
	}

	// 获取树的高度
	int _Height(Node *root)
	{
		if (root == nullptr)
			return 0;

		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);

		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}

	// 左单旋
	void RotateL(Node *parent)
	{
		++_rotateCount;
		Node *subR = parent->_right;
		Node *subRL = subR->_left;

		parent->_right = subRL;
		if (subRL)
			subRL->_parent = parent;

		Node *grandpa = parent->_parent;

		subR->_left = parent;
		parent->_parent = subR;

		if (_root == parent)
		{
			_root = subR;
			subR->_parent = nullptr;
		}
		else
		{
			if (grandpa->_left == parent)
			{
				grandpa->_left = subR;
			}
			else
			{
				grandpa->_right = subR;
			}

			subR->_parent = grandpa;
		}
	}

	// 右单旋
	void RotateR(Node *parent)
	{
		++_rotateCount;
		Node *subL = parent->_left;
		Node *subLR = subL->_right;

		parent->_left = subLR;
		if (subLR)
		{
			subLR->_parent = parent;
		}

		Node *grandpa = parent->_parent;

		subL->_right = parent;
		parent->_parent = subL;

		if (_root == parent)
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			if (grandpa->_left == parent)
			{
				grandpa->_left = subL;
			}
			else
			{
				grandpa->_right = subL;
			}

			subL->_parent = grandpa;
		}
	}

private:
	Node *_root = nullptr;
};

void TestRBTree1()
{
	int a[] = {4, 2, 6, 1, 3, 5, 15, 7, 16, 14};
	int b[] = {16, 3, 7, 11, 9, 26, 18, 14, 15};
	int c[] = {4, 2, 6, 1, 3, 5, 15, 7, 16, 14, 0, 5, 30, 25, 20, 4, 13, 30, 28, 27};
	RBTree<int, int> tree;
	for (auto e : c)
	{
		tree.Insert(make_pair(e, e));
	}

	cout << "此树中序遍历" << endl;
	tree.InOrder();

	if (tree.IsRBTree())
	{
		cout << "此树为红黑树!" << endl;
	}
	else
	{
		cout << "此树不是红黑树!" << endl;
	}
}

void TestRBTree2()
{
	srand(time(0));
	size_t N = 10;
	RBTree<int, int> tree;
	for (size_t i = 0; i < N; ++i)
	{
		int x = rand();
		cout << "Insert:" << x << ":" << i << endl;
		tree.Insert(make_pair(x, i));
	}

	if (tree.IsRBTree())
	{
		cout << "此树为红黑树!" << endl;
	}
	else
	{
		cout << "此树不是红黑树!" << endl;
	}
}
